#!/usr/bin/python3
# -*- coding: utf-8 -*-

import sys
# sys.path.append('/opt/work/caffe/python')
sys.path.insert(0, '.')

import os
import numpy as np
import json

# base_path = r'/rootfs/media/yery/Kaso/data'
base_path = r'/rootfs/media/kasim/Data1/data'

image_base_path = os.path.join(base_path, 'ErisedPersonGender')
json_dir_path = os.path.join(base_path, 'ErisedPersonGender/Annotations.20201105')
# json_dir_path = os.path.join(base_path, 'ErisedPersonGender/Annotations.20201029')
# file_faceid_list_path = os.path.join(base_path, 'ErisedPersonGender/file_faceid_list.txt')
file_faceid_list_path = os.path.join(base_path, 'VideoGenderCropPerson/FaceID/Mtypical_file_faceid_list.txt')

age_group_map = {
    '0-2岁': 0,  # AgeGroup.UNDER_TWO,
    '3-5岁': 1,  # AgeGroup.THREE_TO_FIVE,
    '6-12岁': 2,  # AgeGroup.SIX_TO_TWELVE,
    '13-18岁': 3,  # AgeGroup.THIRTEEN_TO_EIGHTEEN,
    '19-23岁': 4,  # AgeGroup.NINETEEN_TO_TWENTY_THREE,
    '24-29岁': 5,  # AgeGroup.TWENTY_FOUR_TO_TWENTY_NINE,
    '30-35岁': 6,  # AgeGroup.THIRTY_TO_THIRTY_FIVE,
    '36-45岁': 7,  # AgeGroup.THIRTY_SIX_TO_FORTY_FIVE,
    '46-60岁': 8,  # AgeGroup.FORTY_SIX_TO_SIXTY,
    '60岁以上': 9,  # AgeGroup.ABOVE_SIXTY,
    # '跳过': None,  # AgeGroup.UNCERTAIN,
}


def main():
    if not os.path.exists(out_path):
        os.makedirs(out_path)
    os.system('chmod a+wr {}'.format(out_path))

    faceid_files = {}
    with open(file_faceid_list_path, 'r') as file:
        for line in file.readlines():
            # gd/zs/sqqjd/yxxhjt/150100414a54443452067fa6d4be6600/video/20190926_085306/00043_000.jpg 1132
            lines = line.strip().split()
            file_name = lines[0]
            faceid = int(lines[1])
            file_names = faceid_files.get(faceid, None)
            if file_names is None:
                faceid_files[faceid] = [file_name]
            else:
                file_names.append(file_name)

    print('FaceID Count:', len(faceid_files.keys()))

    age_group_count = len(age_group_map.keys())
    age_group_counts_list = []
    total_file_count = 0
    for _, file_names in faceid_files.items():
        file_count = len(file_names)
        total_file_count += file_count
        if file_count > 1:
            device_file_names = {}
            for file_name in file_names:
                device = file_name.split('/')[-4]
                file_names = device_file_names.get(device, None)
                if file_names is None:
                    device_file_names[device] = [file_name]
                else:
                    file_names.append(file_name)

            for device, file_names in device_file_names.items():
                if len(file_names) > 1:
                    age_group_counts = [0] * age_group_count
                    for file_name in file_names:
                        json_file_name = os.path.splitext(file_name.replace('/', '_'))[0] + '.json'
                        json_file_path = os.path.join(json_dir_path, json_file_name)
                        with open(json_file_path, 'r', encoding='utf-8') as f:
                            json_info = json.load(f)
                            if json_info is None:
                                continue
                            name = json_info['outputs']['object'][0]['name']
                            age_group = age_group_map.get(name, None)
                            if age_group is None:
                                continue
                        age_group_counts[age_group] += 1
                    age_group_counts_list.append(age_group_counts)

    age_group_counts = np.array(age_group_counts_list, dtype=np.int32)
    mulit_image_face_count = age_group_counts.shape[0]
    age_group_counts_mask = age_group_counts > 0
    age_group_counts_count = np.sum(age_group_counts_mask, axis=1)
    mask = age_group_counts_count > 1
    age_group_counts = age_group_counts[mask]
    multi_age_group_mulit_image_face_count = age_group_counts.shape[0]
    age_group_total_counts = np.sum(age_group_counts, axis=1)
    age_group_max_counts = np.max(age_group_counts, axis=1)
    age_group_error_counts = age_group_total_counts - age_group_max_counts
    age_group_total_count = age_group_total_counts.sum()
    age_group_error_count = age_group_error_counts.sum()

    idxes = np.arange(0, age_group_counts.shape[0])
    age_group_max_idx0 = np.argmax(age_group_counts, axis=1)
    age_group_max_idx1 = age_group_max_idx0 - 1
    age_group_max_idx2 = age_group_max_idx0 + 1

    start_mask = age_group_max_idx1 < 0
    end_mask = age_group_max_idx2 > (age_group_counts.shape[1]-1)

    age_group_max_idx1[start_mask] = 0
    age_group_max_idx2[end_mask] = 0

    age_group_counts0 = age_group_counts[idxes, age_group_max_idx0]
    age_group_counts1 = age_group_counts[idxes, age_group_max_idx1]
    age_group_counts2 = age_group_counts[idxes, age_group_max_idx2]

    age_group_counts1[start_mask] = 0
    age_group_counts2[end_mask] = 0

    age_group_1off_counts = age_group_counts0 + age_group_counts1 + age_group_counts2
    age_group_error_1off_counts = age_group_total_counts - age_group_1off_counts
    age_group_error_1off_count = age_group_error_1off_counts.sum()

    print('Face Error: {:.02f}% {}, Error: {:.02f}% {}, Total Error: {:.02f}% {}, Total 1Off Error: {:.02f}% {}'.format(
            multi_age_group_mulit_image_face_count*100/mulit_image_face_count, multi_age_group_mulit_image_face_count,
            age_group_error_count*100/age_group_total_count, age_group_error_count,
            age_group_error_count*100/total_file_count, age_group_error_count,
            age_group_error_1off_count*100/total_file_count, age_group_error_1off_count))

    print('Finish!')


if __name__ == '__main__':
    main()
